#include "utils.h"
#include "kpabe_cloud.h"
#include "streamforce_crypto_KPAbe.h"
#include "string.h"
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
#include <pthread.h>
using namespace std;

unordered_map<int,kpabe_pub_t*> cloudPKMap;
unordered_map<int,kpabe_prv_t*> cloudTKMap;
unordered_map<int,element_wrap> cloudDetKeyMap;
pthread_mutex_t lock; 

kpabe_cph_t* convertToCipher(kpabe_pub_t* pubK, unsigned char* cipher, int cipherLen){
	GByteArray *buf =  convertToByteArray(cipher, cipherLen);
	return kpabe_cph_unserialize(pubK, buf, 1);
}

int initACInstance(int streamId, unsigned char* pubBytes, int pubLen, int userId, unsigned char* tkBytes, int tkLen, unsigned char *detK2, int detK2Len){
	kpabe_pub_t* pub = convertToPubKey(pubBytes, pubLen);	
	kpabe_prv_t* prv = convertToPrvKey(pub, tkBytes, tkLen);
	unordered_map<int, kpabe_pub_t*>::const_iterator it = cloudPKMap.find(streamId);
	if (it==cloudPKMap.end())
	{	
		cloudPKMap[streamId] = pub;
	}	
	if (detK2Len!=0){
			element_t tmp;
			element_wrap tp;
			convertToElement(pub, detK2,detK2Len,tmp);
			element_init_Zr(tp.element, pub->p);
			element_set(tp.element,tmp); 
			cloudDetKeyMap[userId] = tp;
	}

	cloudTKMap[userId] = prv;
	return 0;	
}
JNIEXPORT void JNICALL Java_streamforce_crypto_KPAbe_initACInstance
  (JNIEnv *env, jobject obj, jbyteArray pubB, jint streamId, jbyteArray tkB, jint userId,
jbyteArray dkB){
	int pubL = env->GetArrayLength(pubB);
	unsigned char* pubBytes = (unsigned char*)malloc(pubL*sizeof(char));
	env->GetByteArrayRegion(pubB,0,pubL,(jbyte*)pubBytes);

	int prvL = env->GetArrayLength(tkB);
	unsigned char* prvBytes = (unsigned char*)malloc(prvL*sizeof(char));
	env->GetByteArrayRegion(tkB,0,prvL,(jbyte*)prvBytes);

	int detL = env->GetArrayLength(dkB);
	unsigned char *detBytes = (unsigned char*)malloc(detL*sizeof(char));
	env->GetByteArrayRegion(dkB,0,detL,(jbyte*)detBytes);
	
	initACInstance(streamId, pubBytes, pubL, userId, prvBytes, prvL,detBytes,detL);		
}

//transform the ciphertext, using prv key indexed by the policyId
int transform_simple(unsigned char* cipher, int cipherLen, int streamId, int userId, unsigned char** result,
int *len){
	kpabe_pub_t* pub = cloudPKMap[streamId];
	kpabe_prv_t* prv = cloudTKMap[userId];
	kpabe_cph_t* cph = convertToCipher(pub, cipher, cipherLen); 		
	element_t userCipher;
	int s = kpabe_transform(pub,prv,cph,userCipher);
	*len = element_length_in_bytes(userCipher);
	*result = (unsigned char*) malloc((*len)*sizeof(char));
	element_to_bytes(*result, userCipher);
	return s;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_transform_1simple
  (JNIEnv *env, jobject obj, jbyteArray cipher, jint streamId, jint userId){
	int cL = env->GetArrayLength(cipher);
	unsigned char* cipherBytes = (unsigned char*)malloc(cL*sizeof(char));
	env->GetByteArrayRegion(cipher,0,cL,(jbyte*)cipherBytes);
	unsigned char* transformedCipher;
	int transformedCipherLen;
	transform_simple(cipherBytes, cL, streamId, userId, &transformedCipher,
&transformedCipherLen);
	jbyteArray ret = env->NewByteArray(transformedCipherLen);
	env->SetByteArrayRegion(ret,0,transformedCipherLen,(jbyte*)transformedCipher);
	return ret;
}

//multiplying all ciphers
int multiplyCiphers(unsigned char** ciphers, int size, int streamId, unsigned char** result, int*
len){
	kpabe_pub_t* pub = cloudPKMap[streamId];
	element_t one, prod, cipher;
	element_init_GT(one, pub->p);
	element_set_si(one,1);
	element_init_GT(prod, pub->p);
	element_init_GT(cipher, pub->p);
	for (int i=0; i<size; i++){
		element_from_bytes(cipher, ciphers[i]);
		element_mul(prod,cipher,one);
		element_set(one,prod);
	}	
	*len = element_length_in_bytes(prod);
	*result = (unsigned char*) malloc((*len)*sizeof(char));
	element_to_bytes(*result,prod);
	return 0;
}

int divideCiphers(int streamId, unsigned char* cipher1, unsigned char* cipher2, unsigned char** result, int *len){
	kpabe_pub_t* pub = cloudPKMap[streamId];
	element_t c1, c2;
	element_init_GT(c1,pub->p);
	element_init_GT(c2,pub->p);
	element_from_bytes(c1,cipher1);
	element_from_bytes(c2,cipher2);
	
	element_t res;
	element_init_GT(res,pub->p);
	element_div(res,c1,c2);
	*len = element_length_in_bytes(res);
	*result = (unsigned char*)malloc((*len)*sizeof(char));
	element_to_bytes(*result,res);
	return 0;
} 

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_multiplyCiphers
  (JNIEnv *env, jobject obj, jobjectArray ciphers, jint streamId){
	int size = env->GetArrayLength(ciphers);
	unsigned char** cipherArray = (unsigned char**)malloc(size*sizeof(char*));
	for (int i=0; i<size; i++){
		jbyteArray tmp = (jbyteArray)env->GetObjectArrayElement(ciphers,i);
		int len = env->GetArrayLength(tmp);
		cipherArray[i] = (unsigned char*)malloc(len*sizeof(char));
		env->GetByteArrayRegion(tmp,0,len,(jbyte*)cipherArray[i]);
	}

	unsigned char *result;
	int resultLen;
	multiplyCiphers(cipherArray,size,streamId,&result,&resultLen);
	
	jbyteArray ret = env->NewByteArray(resultLen);
	env->SetByteArrayRegion(ret,0,resultLen,(jbyte*)result);
	return ret;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_divideCiphers
  (JNIEnv *env, jobject obj, jbyteArray cipher1, jbyteArray cipher2, jint streamId){
	int size = env->GetArrayLength(cipher1);
	unsigned char *cipherArray1 = (unsigned char*)malloc(size*sizeof(char));
	unsigned char *cipherArray2 = (unsigned char*)malloc(size*sizeof(char));
	env->GetByteArrayRegion(cipher1,0,size,(jbyte*)cipherArray1);
	env->GetByteArrayRegion(cipher2,0,size,(jbyte*)cipherArray2);	

	unsigned char *result;
	int len;
	divideCiphers(streamId,cipherArray1,cipherArray2,&result,&len);
	jbyteArray ret = env->NewByteArray(len);
	env->SetByteArrayRegion(ret,0,len,(jbyte*)result);
	return 0;
}



int transform_join(int streamId, int userId, unsigned char *cipher, int cipherLen, unsigned char **transformedCipher, int
*transformedCipherLen){
	element_t tpm;
	transform_join(streamId, userId, cipher, cipherLen,tpm);
	*transformedCipherLen = element_length_in_bytes(tpm); 
	*transformedCipher = (unsigned char*) malloc(*transformedCipherLen*sizeof(char));
	element_to_bytes(*transformedCipher, tpm); 
	return 0; 	
}

//the cipher is in Zr
int transform_join(int streamId, int userId, unsigned char *cipher, int cipherLen, element_t cph){
	kpabe_pub_t* pub = cloudPKMap[streamId]; 
	element_t tmp;
	element_init_G1(tmp, pub->p);
	element_from_bytes(tmp, cipher);
	
	element_init_G1(cph, pub->p);
	element_pow_zn(cph,tmp,cloudDetKeyMap[userId].element);
	
	return 0; 	
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_transform_1join
  (JNIEnv *env, jobject obj, jint streamId, jint userId, jbyteArray cipher){
	int len = env->GetArrayLength(cipher);
	unsigned char *cipherBytes = (unsigned char*)malloc(len*sizeof(char));
	env->GetByteArrayRegion(cipher,0,len,(jbyte*)cipherBytes);
	unsigned char *returnedBytes;
	int transformLen;
	transform_join(streamId,userId,cipherBytes,len,&returnedBytes,&transformLen);
	jbyteArray ret = env->NewByteArray(transformLen);
	env->SetByteArrayRegion(ret,0,transformLen,(jbyte*)returnedBytes);
	return ret;
}
